/**
 * 
 */
package jmine.tec.utils.cripto.cipher;

import java.io.DataOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.security.InvalidKeyException;
import java.security.Key;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;

/**
 * @author takeshi
 * @version 20/08/2007
 */
public class CipherOutputStream extends OutputStream {

    /**
     * 
     */
    private static final int DEFAULT_BUFFER_SIZE = 512;

    private byte[] buffer;

    private final Cipher cipher;

    private boolean closed = false;

    private final DataOutputStream output;

    private int written;

    private final int chunkSize;

    /**
     * C'tor
     * 
     * @param output the {@link DataOutputStream}
     * @param cipher the cipher to encode
     * @param key the key to encode
     * @throws InvalidKeyException if the key is invalid
     * @param chunkSize the chunk size
     */
    public CipherOutputStream(final OutputStream output, final Cipher cipher, final Key key, final int chunkSize)
            throws InvalidKeyException {
        super();
        this.output = new DataOutputStream(output);
        this.buffer = new byte[DEFAULT_BUFFER_SIZE];
        this.cipher = cipher;
        this.cipher.init(Cipher.ENCRYPT_MODE, key);
        this.chunkSize = chunkSize;
    }

    /**
     * @throws IOException e
     * @see java.io.OutputStream#close()
     */
    @Override
    public synchronized void close() throws IOException {
        try {
            this.flush();
        } finally {
            this.closed = true;
            this.output.close();
        }
    }

    /**
     * @throws IOException e
     * @see java.io.OutputStream#flush()
     */
    @Override
    public synchronized void flush() throws IOException {
        this.encode();
        this.output.flush();
    }

    /**
     * Writes a chunk of bytes into the stream.
     * 
     * @param b byte array to write
     * @param off offset from witch to write
     * @param len amount of bytes to write
     * @throws IOException if an {@link IOException} occurs
     * @see java.io.OutputStream#write(byte[], int, int)
     */
    @Override
    public synchronized void write(final byte[] b, final int off, final int len) throws IOException {
        if (this.closed) {
            throw new IOException("Stream closed");
        }
        if (len == 0) {
            return;
        }
        int copied = 0;
        while (copied < len) {
            int toCopy = Math.min(len - copied, this.availableSpace());
            System.arraycopy(b, off + copied, this.buffer, this.written, toCopy);
            this.written += toCopy;
            copied += toCopy;
            this.encodeAllChunks();
        }
    }

    /**
     * @param b int
     * @throws IOException e
     * @see java.io.OutputStream#write(int)
     */
    @Override
    public synchronized void write(final int b) throws IOException {
        if (this.closed) {
            throw new IOException("Stream closed");
        }
        this.buffer[this.written++] = (byte) b;
        if (this.written >= this.chunkSize) {
            this.encode();
        }
    }

    /**
     * @throws IOException if an {@link IOException} occurs
     */
    private void encodeAllChunks() throws IOException {
        while (this.written >= this.chunkSize) {
            this.encode();
        }
    }

    /**
     * @return int the amount of space in bytes available
     */
    private int availableSpace() {
        return this.buffer.length - this.written;
    }

    /**
     * @throws IOException if an {@link IOException} if thrown
     */
    private void encode() throws IOException {
        if (this.written <= 0) { // nothing to write
            return;
        }
        int amountToEncode = Math.min(this.chunkSize, this.written);

        byte[] encoded;
        try {
            encoded = this.cipher.doFinal(this.buffer, 0, amountToEncode);
        } catch (IllegalBlockSizeException e) {
            throw new IllegalStateException(e);
        } catch (BadPaddingException e) {
            throw new IllegalStateException(e);
        }
        this.output.writeInt(encoded.length);
        this.output.write(encoded);

        // move bytes around
        System.arraycopy(this.buffer, amountToEncode, this.buffer, 0, this.buffer.length - amountToEncode);

        this.written -= amountToEncode;
    }

}
